from .base_stats import *
from scipy.interpolate import make_interp_spline, CubicSpline
class DistributeFit:
	'''
	拟合分布类
	'''
	def __init__(self):
		self.dist: Distribution = None # type: ignore # 待拟合分布

	def plot_pdf(self, x_options: tuple[float, float] = None, y_options: tuple[float, float] = None, color: ColorType = None, label: str = '', count: int = 1000)->tuple[FloatOrArray, FloatOrArray]	: # type: ignore
		'''
		绘制分布曲线
		'''
		x_options = x_options or (0, 1)
		y_options = y_options or (0, 1)
		x_min, x_max = self.dist.x_range
		x = np.linspace(x_min, x_max, count) # type: ignore
		x, y = self._calc_pdf_for_plot(x, x_options, y_options)
		plt.plot(x, y, '-', color=color, label=label)
		return x, y
	
	def bind(self, dist: Distribution):
		'''
		绑定分布
		'''
		self.dist = dist

	def pdf(self, x: FloatOrArray)->FloatOrArray: # type: ignore
		'''
		计算概率密度
		params:
			x: list[float], 待计算的x值
		returns:
			y: list[float], 概率密度
		'''
		return StatsPublic.calc_on_array_or_single(lambda x: self.dist.prob if self.dist.in_range(x) else 0, x)

	def _calc_pdf_for_plot(self, x: FloatOrArray, x_options: tuple[float, float] = None, y_options: tuple[float, float] = None)->tuple[FloatOrArray, FloatOrArray]: # type: ignore
		'''
		概率密度函数
		params:
			x: list[float], 待计算的x值
			x_options: tuple[float, float], x轴范围
			y_options: tuple[float, float], y轴范围
		returns:
			x: list[float], 绘图的x值
			y: list[list[float]], 绘图的y值
		'''
		y = self.pdf(x)
		x =  StatsPublic.scale_and_offset(x, x_options)
		y =  StatsPublic.scale_and_offset(y, y_options)
		return x,y

class DistFitGroups(DistributeFit):
	'''组合分布拟合类'''
	dist: GroupedDistributions # type: ignore
	dist_fits: list[DistributeFit]
	def __init__(self, fit_class: type[DistributeFit]):
		super().__init__()
		self.fit_class = fit_class
		self.dist_fits: list[DistributeFit] = []

	def bind(self, dist: Distribution):
		super().bind(dist)
		for dist in self.dist.dists:
			dist_fit = self.fit_class()
			dist_fit.bind(dist)
			self.dist_fits.append(dist_fit)

	def pdf(self, x: FloatOrArray):
		def pdf_sum(x):
			return sum(dist_fit.pdf(x) for dist_fit in self.dist_fits)
		return StatsPublic.calc_on_array_or_single(pdf_sum, x)

	def plot_pdf(self, x_options: tuple[float, float] = None, y_options: tuple[float, float] = None, color: ColorType = None, label: str = '', show_groups: bool = False, count: int = 1000)->tuple[FloatOrArray, FloatOrArray]	: # type: ignore
		'''
		绘制分布曲线
		'''
		x, y = super().plot_pdf(x_options, y_options, color, f'{label}总体分布', count)
		if show_groups:
			for i in range(len(self.dist_fits)):
				dist_fit_i = self.dist_fits[i]
				dist_i: Distribution = dist_fit_i.dist
				plt.axvline(dist_i.lower_bound, color='gray', linestyle='--')
				plt.text(dist_i.lower_bound + self.dist.group_size/2, 0.9, f'组{i+1}', ha='center', va='top')
				print(f'第{i}组: e0:{dist_i.expect0:.3f}, e1:{dist_i.expect1:.3f}, e2:{dist_i.expect2:.3f}, e3:{dist_i.expect3:.3f}, e4:{dist_i.expect4:.3f}')
				dist_fit_i.plot_pdf(x_options=x_options, y_options=y_options, label=f'组{i+1}')
		return x, y

	def plot_median_pdf(self, x_options: tuple[float, float] = None, y_options: tuple[float, float] = None, color: ColorType = None, label: str = '', interpolate_count: int = 0)->tuple[FloatOrArray, list[FloatOrArray]]	: # type: ignore
		'''
		绘制中位数分布曲线
		'''
		x = np.array([StatsPublic.scale_and_offset(dist.median, x_options) for dist in self.dist.dists])
		y = np.array([StatsPublic.scale_and_offset(dist.prob, y_options) for dist in self.dist.dists])
		if interpolate_count > 0:
			# 生成三次样条插值函数
			spline = make_interp_spline(x, y, k=3)  # k=3表示三次样条
			x_smooth = np.linspace(np.min(x), np.max(x), interpolate_count) 
			y_smooth = spline(x_smooth)
			plt.plot(x_smooth, y_smooth, color=color, label=label)
			return x_smooth , [y_smooth]
		else:
			plt.plot(x, y, color=color, label=label)

			return x, [y]

class DistFitGroupsCubicSpline(DistFitGroups):
	'''三次样条组合分布拟合类'''
	def __init__(self):
		super().__init__(DistributeFit)
		self.spline: CubicSpline = None # type: ignore

	def bind(self, dist: Distribution):
		super().bind(dist)
		x = [dist.median for dist in self.dist.dists]
		y = [dist.prob for dist in self.dist.dists]
		self.spline = CubicSpline(x, y, bc_type='natural')

	def pdf(self, x: FloatOrArray):
		def calc_pdf(x):
			return self.spline(x)
		return StatsPublic.calc_on_array_or_single(calc_pdf, x)

